package plaindb

import (
	"encoding/json"
	"fhe-search-cloud/g"
	"fhe-search-cloud/utils"
	"fhe-search-cloud/utils/log"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

type SrcDocument struct {
	Id    string `json:"id"`
	Title string `json:"title"`
	Text  string `json:"text"`
}
type DestDocument struct {
	Id       int    `db:"id"`
	Title    string `db:"title"`
	FileType string `db:"file_type"`
	FilePath string `db:"file_path"`
}

func DataSetLoadToDB() {
	// 数据集来源
	// https://github.com/brightmart/nlp_chinese_corpus
	// https://drive.google.com/file/d/1EdHUZIDpgcBoSqbjlfNKJ3b1t0XIUjbt/view?usp=sharing
	log.Log("start load data to db.")
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/mysql",
		g.GetConfig().MysqlConfig.User,
		g.GetConfig().MysqlConfig.Password,
		g.GetConfig().MysqlConfig.Host,
		g.GetConfig().MysqlConfig.Port,
	)
	destDb, err := sqlx.Connect("mysql", dsn)
	if err != nil {
		panic(err)
	}
	defer destDb.Close()
	fileBasePath := g.GetConfig().FheSearchConfig.FileBasePath
	srcDocRootPath := fileBasePath + "wiki_zh/"
	destDocDir := "AllPlainData/"
	destDocRootPath := fileBasePath + destDocDir

	// init dest db
	var sql string
	log.Log("start clear dest db.")
	sql = `CREATE DATABASE IF NOT EXISTS wiki2019zh;`
	destDb.MustExec(sql)
	destDb.MustExec(`USE wiki2019zh;`)
	sql = `DROP TABLE IF EXISTS documents;`
	destDb.MustExec(sql)
	sql = `CREATE TABLE IF NOT EXISTS documents (
		id BIGINT PRIMARY KEY AUTO_INCREMENT,
		title VARCHAR(255) NOT NULL,
		file_type VARCHAR(255) NOT NULL,
		file_path VARCHAR(255) NOT NULL);`
	destDb.MustExec(sql)
	log.Log("init dest db finished.")
	// init dest dir
	log.Log("start clear dest dir.")
	utils.ClearDir(destDocRootPath)
	log.Log("init dest dir finished.")

	const ChanCap = g.MysqlMaxAllUnitNum / 4 * 2
	destDocChan := make(chan *DestDocument, ChanCap)
	wg := &sync.WaitGroup{}
	wg.Add(1)

	// write dest db 消费destDocChan
	// 批量向数据库中插入数据
	go func() {
		defer wg.Done()
		sql := "insert into documents (id, title, file_type, file_path) values (:id, :title, :file_type, :file_path)"
		docItems := []interface{}{}
		for {
			docItem := <-destDocChan
			if docItem == nil {
				break
			}
			docItems = append(docItems, docItem)
			if len(docItems) >= ChanCap/2 {
				destDb.NamedExec(sql, docItems)
				docItems = []interface{}{}
			}
		}
		if len(docItems) > 0 {
			destDb.NamedExec(sql, docItems)
		}
	}()

	// read src docs 向destDocChan写入数据
	// 遍历wiki_zh2019数据集，将每个文档写入一个单独的文件
	// 并且在数据库wiki2019zh/documents中添加记录
	cnt := 0
	err = filepath.Walk(srcDocRootPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		fileContent, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		lines := strings.Split(string(fileContent), "\n")
		for _, line := range lines {
			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}
			doc := &SrcDocument{}
			err := json.Unmarshal([]byte(line), doc)
			if err != nil {
				panic(err)
			}
			destPath := destDocDir + doc.Id + ".txt"
			idInt, err := strconv.Atoi(doc.Id)
			if err != nil {
				panic(err)
			}
			destDoc := &DestDocument{
				Id:       idInt,
				Title:    doc.Title,
				FileType: "txt",
				FilePath: destPath,
			}
			destDocChan <- destDoc
			// write dest doc
			destDocPath := fileBasePath + destPath
			err = os.WriteFile(destDocPath, []byte(doc.Text), os.ModePerm)
			if err != nil {
				panic(err)
			}
			cnt++
			if cnt%10000 == 0 {
				log.Log("doc cnt:", cnt, "finish load.")
			}
		}
		return nil
	})
	if err != nil {
		panic(err)
	}
	close(destDocChan)
	wg.Wait()
	log.Log("doc cnt:", cnt)
	log.Log("load data to db done.")
}
